From 8c4cb846ba1634f1e149833090beb7376f216157 Mon Sep 17 00:00:00 2001
From: Sebastian Rasmussen <sebras@gmail.com>
Date: Sun, 1 Aug 2021 20:07:14 +0200
Subject: [PATCH] jbig2dec: Improve reporting from custom allocator that limits
 memory use.

---
 jbig2dec.c | 65 +++++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 45 insertions(+), 20 deletions(-)

diff --git a/jbig2dec.c b/jbig2dec.c
index 3996e5b..dc1fd56 100644
--- a/jbig2dec.c
+++ b/jbig2dec.c
@@ -91,6 +91,47 @@ static int print_usage(void);
 #define KBYTE 1024
 #define MBYTE (1024 * KBYTE)
 
+static void *jbig2dec_reached_limit(jbig2dec_allocator_t *allocator, size_t oldsize, size_t size)
+{
+    size_t limit_mb = allocator->memory_limit / MBYTE;
+    size_t used_mb = allocator->memory_used / MBYTE;
+    size_t oldsize_mb = oldsize / MBYTE;
+    size_t size_mb = size / MBYTE;
+
+    if (oldsize == 0)
+        jbig2_error(allocator->ctx, JBIG2_SEVERITY_FATAL, -1, "memory: limit reached: limit: %zu (%zu Mbyte) used: %zu (%zu Mbyte) allocation: %zu (%zu Mbyte)",
+            allocator->memory_limit, limit_mb,
+            allocator->memory_used, used_mb,
+            size, size_mb);
+    else
+        jbig2_error(allocator->ctx, JBIG2_SEVERITY_FATAL, -1, "memory: limit reached: limit: %zu (%zu Mbyte) used: %zu (%zu Mbyte) reallocation: %zu (%zu Mbyte) -> %zu (%zu Mbyte)",
+            allocator->memory_limit, limit_mb,
+            allocator->memory_used, used_mb,
+            oldsize, oldsize_mb,
+            size, size_mb);
+
+    return NULL;
+}
+
+static void jbig2dec_peak(jbig2dec_allocator_t *allocator)
+{
+    size_t limit_mb = allocator->memory_limit / MBYTE;
+    size_t peak_mb = allocator->memory_peak / MBYTE;
+    size_t used_mb = allocator->memory_used / MBYTE;
+
+    if (allocator->ctx == NULL)
+        return;
+    if (used_mb <= peak_mb)
+        return;
+
+    allocator->memory_peak = allocator->memory_used;
+
+    jbig2_error(allocator->ctx, JBIG2_SEVERITY_DEBUG, JBIG2_UNKNOWN_SEGMENT_NUMBER, "memory: limit: %lu %sbyte used: %lu %sbyte, peak: %lu %sbyte",
+        limit_mb > 0 ? limit_mb : allocator->memory_limit, limit_mb > 0 ? "M" : "",
+        used_mb > 0 ? used_mb : allocator->memory_used, used_mb > 0 ? "M" : "",
+        peak_mb > 0 ? peak_mb : allocator->memory_peak, peak_mb > 0 ? "M" : "");
+}
+
 static void *jbig2dec_alloc(Jbig2Allocator *allocator_, size_t size)
 {
     jbig2dec_allocator_t *allocator = (jbig2dec_allocator_t *) allocator_;
@@ -102,7 +143,7 @@ static void *jbig2dec_alloc(Jbig2Allocator *allocator_, size_t size)
         return NULL;
 
     if (size + ALIGNMENT > allocator->memory_limit - allocator->memory_used)
-        return NULL;
+        return jbig2dec_reached_limit(allocator, 0, size + ALIGNMENT);
 
     ptr = malloc(size + ALIGNMENT);
     if (ptr == NULL)
@@ -110,15 +151,7 @@ static void *jbig2dec_alloc(Jbig2Allocator *allocator_, size_t size)
     memcpy(ptr, &size, sizeof(size));
     allocator->memory_used += size + ALIGNMENT;
 
-    if (allocator->memory_used > allocator->memory_peak) {
-        allocator->memory_peak = allocator->memory_used;
-
-        if (allocator->ctx) {
-            size_t limit_mb = allocator->memory_limit / MBYTE;
-            size_t peak_mb = allocator->memory_peak / MBYTE;
-            jbig2_error(allocator->ctx, JBIG2_SEVERITY_DEBUG, JBIG2_UNKNOWN_SEGMENT_NUMBER, "memory: limit: %lu Mbyte peak usage: %lu Mbyte", limit_mb, peak_mb);
-        }
-    }
+    jbig2dec_peak(allocator);
 
     return (unsigned char *) ptr + ALIGNMENT;
 }
@@ -158,7 +191,7 @@ static void *jbig2dec_realloc(Jbig2Allocator *allocator_, void *p, size_t size)
     memcpy(&oldsize, oldp, sizeof(oldsize));
 
     if (size + ALIGNMENT > allocator->memory_limit - allocator->memory_used + oldsize + ALIGNMENT)
-        return NULL;
+        return jbig2dec_reached_limit(allocator, oldsize + ALIGNMENT, size + ALIGNMENT);
 
     p = realloc(oldp, size + ALIGNMENT);
     if (p == NULL)
@@ -168,15 +201,7 @@ static void *jbig2dec_realloc(Jbig2Allocator *allocator_, void *p, size_t size)
     memcpy(p, &size, sizeof(size));
     allocator->memory_used += size + ALIGNMENT;
 
-    if (allocator->memory_used > allocator->memory_peak) {
-        allocator->memory_peak = allocator->memory_used;
-
-        if (allocator->ctx) {
-            size_t limit_mb = allocator->memory_limit / MBYTE;
-            size_t peak_mb = allocator->memory_peak / MBYTE;
-            jbig2_error(allocator->ctx, JBIG2_SEVERITY_DEBUG, JBIG2_UNKNOWN_SEGMENT_NUMBER, "memory: limit: %lu Mbyte peak usage: %lu Mbyte", limit_mb, peak_mb);
-        }
-    }
+    jbig2dec_peak(allocator);
 
     return (unsigned char *) p + ALIGNMENT;
 }
